from typing import Tuple, List, Set, Dict, Any, Optional
from dataclasses import dataclass


@dataclass
class Constraint:
    """
    CBS约束类，表示在特定时间步不能到达的位置
    """
    position: Tuple[int, int]  # 位置 (x, y)
    timestep: int  # 时间步
    agent_id: int  # 约束适用的机器人ID，-1表示适用于所有机器人
    
    def __str__(self) -> str:
        """字符串表示"""
        agent_str = f"Agent {self.agent_id}" if self.agent_id >= 0 else "All agents"
        return f"{agent_str} cannot be at position {self.position} at timestep {self.timestep}"
    
    def __eq__(self, other) -> bool:
        """判断两个约束是否相等"""
        if not isinstance(other, Constraint):
            return False
        return (self.position == other.position and 
                self.timestep == other.timestep and 
                self.agent_id == other.agent_id)
    
    def __hash__(self) -> int:
        """哈希函数，用于在集合中去重"""
        return hash((self.position, self.timestep, self.agent_id))


@dataclass
class EdgeConstraint:
    """
    边约束，表示在特定时间步不能从一个位置移动到另一个位置
    """
    position_a: Tuple[int, int]  # 起始位置 (x, y)
    position_b: Tuple[int, int]  # 目标位置 (x, y)
    timestep: int  # 时间步
    agent_id: int  # 约束适用的机器人ID，-1表示适用于所有机器人
    
    def __str__(self) -> str:
        """字符串表示"""
        agent_str = f"Agent {self.agent_id}" if self.agent_id >= 0 else "All agents"
        return f"{agent_str} cannot move from {self.position_a} to {self.position_b} at timestep {self.timestep}"
    
    def __eq__(self, other) -> bool:
        """判断两个边约束是否相等"""
        if not isinstance(other, EdgeConstraint):
            return False
        return (self.position_a == other.position_a and 
                self.position_b == other.position_b and 
                self.timestep == other.timestep and 
                self.agent_id == other.agent_id)
    
    def __hash__(self) -> int:
        """哈希函数，用于在集合中去重"""
        return hash((self.position_a, self.position_b, self.timestep, self.agent_id))


class ConstraintSet:
    """
    约束集合，包含多个约束
    """
    def __init__(self):
        """初始化空约束集合"""
        self.vertex_constraints: Set[Constraint] = set()  # 顶点约束
        self.edge_constraints: Set[EdgeConstraint] = set()  # 边约束
    
    def add_constraint(self, constraint: Constraint):
        """添加顶点约束"""
        self.vertex_constraints.add(constraint)
    
    def add_edge_constraint(self, edge_constraint: EdgeConstraint):
        """添加边约束"""
        self.edge_constraints.add(edge_constraint)
    
    def get_constraints(self, agent_id: int) -> Set[Tuple[int, int, int]]:
        """
        获取指定机器人的所有顶点约束
        
        Args:
            agent_id: 机器人ID
            
        Returns:
            Set[Tuple[int, int, int]]: 约束集合 (x, y, timestep)
        """
        constraints = set()
        for constraint in self.vertex_constraints:
            if constraint.agent_id == agent_id or constraint.agent_id == -1:
                constraints.add((constraint.position[0], constraint.position[1], constraint.timestep))
        return constraints
    
    def get_edge_constraints(self, agent_id: int) -> Set[Tuple[Tuple[int, int], Tuple[int, int], int]]:
        """
        获取指定机器人的所有边约束
        
        Args:
            agent_id: 机器人ID
            
        Returns:
            Set[Tuple[Tuple[int, int], Tuple[int, int], int]]: 边约束集合 ((x1, y1), (x2, y2), timestep)
        """
        constraints = set()
        for constraint in self.edge_constraints:
            if constraint.agent_id == agent_id or constraint.agent_id == -1:
                constraints.add((constraint.position_a, constraint.position_b, constraint.timestep))
        return constraints
    
    def __str__(self) -> str:
        """字符串表示"""
        result = "Vertex Constraints:\n"
        for constraint in self.vertex_constraints:
            result += f"  {str(constraint)}\n"
        result += "Edge Constraints:\n"
        for constraint in self.edge_constraints:
            result += f"  {str(constraint)}\n"
        return result 